{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Yolo v2 介绍\n",
    "## 改进策略\n",
    "### 1. Batch Normalization\n",
    "Batch Normalization可以提升模型收敛速度，而且可以起到一定正则化效果，降低模型的过拟合。在YOLOv2中，每个卷积层后面都添加了Batch Normalization层，并且不再使用droput。使用Batch Normalization后，YOLOv2的mAP提升了2.4%。\n",
    "### 2. 高分辨率\n",
    "目前大部分的检测模型都会在先在ImageNet分类数据集上预训练模型的主体部分（CNN特征提取器），由于历史原因，ImageNet分类模型基本采用大小为224x224的图片作为输入，分辨率相对较低，不利于检测模型。所以YOLOv1在采用224x224分类模型预训练后，将分辨率增加至448x448 ，并使用这个高分辨率在检测数据集上finetune。但是直接切换分辨率，检测模型可能难以快速适应高分辨率。所以**YOLOv2增加了在ImageNet数据集上使用448x448输入来finetune分类网络这一中间过程（10 epochs），这可以使得模型在检测数据集上finetune之前已经适用高分辨率输入。**使用高分辨率分类器后，YOLOv2的mAP提升了约4%。\n",
    "### 3. 带Anchor Boxes的卷积\n",
    "### 4. 聚类"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import cv2\n",
    "import torchvision.transforms as transforms\n",
    "from torch.autograd import Variable"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 数据集相关"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "source": [
    "class Config:\n",
    "    inp_size = np.array([416, 416], dtype=np.int)   #输入图片尺寸\n",
    "    out_size = inp_size / 32                        #输出网络特征图尺寸\n",
    "    out_size = np.array(out_size, dtype=np.int)\n",
    "    thresh = 0.5  #Pr(Object)*Pr(Classi|Object)\n",
    "    nms_thresh = 0.3\n",
    "    \n",
    "    def __init__(self) -> None:\n",
    "        pass\n",
    "\n",
    "# VOC\n",
    "class VOC:\n",
    "    def __init__(self) -> None:\n",
    "        self.label_names = ('aeroplane', 'bicycle', 'bird', 'boat',\n",
    "                'bottle', 'bus', 'car', 'cat', 'chair',\n",
    "                'cow', 'diningtable', 'dog', 'horse',\n",
    "                'motorbike', 'person', 'pottedplant',\n",
    "                'sheep', 'sofa', 'train', 'tvmonitor')\n",
    "        self.num_classes = len(self.label_names)\n",
    "        # pw, ph\n",
    "        self.anchors = np.asarray([(1.08, 1.19), (3.42, 4.41),\n",
    "                        (6.63, 11.38), (9.42, 5.11), (16.62, 10.52)],\n",
    "                        dtype=np.float)\n",
    "        self.num_anchors = len(self.anchors)\n",
    "        base = int(np.ceil(pow(self.num_classes, 1. / 3)))\n",
    "        # self.colors = [self._to_color(x, base) for x in range(self.num_classes)]\n",
    "        self.colors = [(np.random.randint(255),\n",
    "                     np.random.randint(255),\n",
    "                     np.random.randint(255)) for _ in range(self.num_classes)]\n",
    "    def _to_color( self, indx, base):\n",
    "        base2 = base * base\n",
    "        b = 2 - indx / base2\n",
    "        r = 2 - (indx % base2) / base\n",
    "        g = 2 - (indx % base2) % base\n",
    "        return b * 127, r * 127, g * 127\n",
    "voc = VOC()"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 网络结构"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## 基础模块"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "source": [
    "class Conv2d(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels, kernel_size, stride=1,\n",
    "                 relu=True, same_padding=False):\n",
    "        super(Conv2d, self).__init__()\n",
    "        padding = int((kernel_size - 1) / 2) if same_padding else 0\n",
    "        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size,\n",
    "                              stride, padding=padding)\n",
    "        self.relu = nn.LeakyReLU(0.1, inplace=True) if relu else None\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv(x)\n",
    "        if self.relu is not None:\n",
    "            x = self.relu(x)\n",
    "        return x\n",
    "\n",
    "class Conv2d_BatchNorm(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels, kernel_size, stride=1,\n",
    "                 relu=True, same_padding=False):\n",
    "        super(Conv2d_BatchNorm, self).__init__()\n",
    "        padding = int((kernel_size-1) / 2) if same_padding else 0   #当stride=1时，通过padding使得图片大小不改变\n",
    "\n",
    "        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size,\n",
    "                              stride, padding=padding, bias=False)\n",
    "        self.bn = nn.BatchNorm2d(out_channels, momentum=0.01)\n",
    "        self.relu = nn.LeakyReLU(0.1, inplace=True) if relu else None\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.conv(x)\n",
    "        x = self.bn(x)\n",
    "        if self.relu is not None:\n",
    "            x = self.relu(x)\n",
    "        return x\n",
    "\n",
    "class FC(nn.Module):\n",
    "    def __init__(self, in_features, out_features, relu=True):\n",
    "        super(FC, self).__init__()\n",
    "        self.fc = nn.Linear(in_features, out_features)\n",
    "        self.relu = nn.ReLU(inplace=True) if relu else None\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.fc(x)\n",
    "        if self.relu is not None:\n",
    "            x = self.relu(x)\n",
    "        return x\n",
    "\n",
    "class reorg_layer(nn.Module):\n",
    "    def __init__(self, stride):\n",
    "        super().__init__()\n",
    "        self.stride = stride\n",
    "    def forward(self, x):\n",
    "        batch_size, channels, height, width = x.size()\n",
    "        _height, _width = height // self.stride, width // self.stride\n",
    "        x = x.view(batch_size, channels, _height, self.stride, _width, self.stride).transpose(3, 4).contiguous()\n",
    "        x = x.view(batch_size, channels, _height * _width, self.stride * self.stride).transpose(2, 3).contiguous()\n",
    "        x = x.view(batch_size, channels, self.stride * self.stride, _height, _width).transpose(1, 2).contiguous()\n",
    "        x = x.view(batch_size, -1, _height, _width)\n",
    "\n",
    "        return x\n"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Darknet19"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "source": [
    "class Darknet19(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Darknet19, self).__init__()\n",
    "        # [输出通道数，kernel_szie]\n",
    "        net_cfgs = [\n",
    "            # conv1s--0-4\n",
    "            [(32, 3)],\n",
    "            ['M', (64, 3)],\n",
    "            ['M', (128, 3), (64, 1), (128, 3)],\n",
    "            ['M', (256, 3), (128, 1), (256, 3)],\n",
    "            ['M', (512, 3), (256, 1), (512, 3), (256, 1), (512, 3)],\n",
    "            # conv2--5,对应self.conv6\n",
    "            ['M', (1024, 3), (512, 1), (1024, 3), (512, 1), (1024, 3)],\n",
    "            # ------------\n",
    "            # conv3--6\n",
    "            [(1024, 3), (1024, 3)],\n",
    "            # conv4--7\n",
    "            [(1024, 3)]\n",
    "        ]\n",
    "        #darknet\n",
    "        self.conv1s, c1 = self._make_layers(3, net_cfgs[0:5])\n",
    "        self.conv2, c2 = self._make_layers(c1, net_cfgs[5])\n",
    "        self.conv3, c3 = self._make_layers(c2, net_cfgs[6])\n",
    "        stride=2\n",
    "        self.reorg = reorg_layer(stride=2)\n",
    "        self.conv4, c4 = self._make_layers((c1*(stride*stride) + c3), net_cfgs[7])\n",
    "\n",
    "        # 线性层\n",
    "        out_channels = voc.num_anchors * (voc.num_classes + 5)\n",
    "        self.conv5 = Conv2d(c4, out_channels, 1, 1, relu=False)\n",
    "        self.global_average_pool = nn.AvgPool2d((1, 1))\n",
    "\n",
    "    def _make_layers(self, in_channels, net_cfg):\n",
    "        layers = []\n",
    "\n",
    "        if len(net_cfg) > 0 and isinstance(net_cfg[0], list):\n",
    "            for sub_cfg in net_cfg:\n",
    "                layer, in_channels = self._make_layers(in_channels, sub_cfg)\n",
    "                layers.append(layer)\n",
    "        else:\n",
    "            for item in net_cfg:\n",
    "                if item == 'M':\n",
    "                    layers.append(nn.MaxPool2d(kernel_size=2, stride=2))\n",
    "                else:\n",
    "                    out_channels, ksize = item\n",
    "                    layers.append(Conv2d_BatchNorm(in_channels, out_channels, ksize,same_padding=True))\n",
    "                    in_channels = out_channels\n",
    "        return nn.Sequential(*layers), in_channels\n",
    "    \n",
    "    def forward(self, im_data, gt_boxes=None, gt_classes=None, dontcare=None,\n",
    "                size_index=0):\n",
    "        conv1s = self.conv1s(im_data)\n",
    "        conv2 = self.conv2(conv1s)\n",
    "        conv3 = self.conv3(conv2)\n",
    "        conv1s_reorg = self.reorg(conv1s)\n",
    "        cat_1_3 = torch.cat([conv1s_reorg, conv3], 1)\n",
    "        conv4 = self.conv4(cat_1_3)\n",
    "        conv5 = self.conv5(conv4)   # batch_size, out_channels, h, w\n",
    "        global_average_pool = self.global_average_pool(conv5)   # 1x1的平均池化有什么用\n",
    "\n",
    "        bsize, _, h, w = global_average_pool.size()\n",
    "        # [B, num_anchor*25, H, W]->[B, H, W, num_anchor*25]-->[B, H*W, num_anchor, 25]\n",
    "        global_average_pool_reshaped = \\\n",
    "            global_average_pool.permute(0, 2, 3, 1).contiguous().view(bsize, -1, voc.num_anchors, voc.num_classes + 5)\n",
    "        # sigmoid使得xy小于1，偏移之后始终位于cell之内\n",
    "        xy_pred = F.sigmoid(global_average_pool_reshaped[:, :, :, 0:2])\n",
    "        wh_pred = torch.exp(global_average_pool_reshaped[:, :, :, 2:4])\n",
    "        bbox_pred = torch.cat([xy_pred, wh_pred], 3)\n",
    "        iou_pred = F.sigmoid(global_average_pool_reshaped[:, :, :, 4:5])\n",
    "        score_pred = global_average_pool_reshaped[:, :, :, 5:].contiguous()\n",
    "        # 通过soft_max求类别条件概率\n",
    "        prob_pred = F.softmax(score_pred.view(-1, score_pred.size()[-1])).view_as(score_pred)\n",
    "\n",
    "        return bbox_pred, iou_pred, prob_pred\n",
    "\n"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "source": [
    "class YOLOV2:\n",
    "    def __init__(self, path_ckpt=None):\n",
    "        self.darknet19 = Darknet19()\n",
    "        if path_ckpt:\n",
    "            import h5py\n",
    "            print('Use ckpt from:',path_ckpt)\n",
    "            h5f = h5py.File(path_ckpt, mode='r')\n",
    "            for k, v in list(self.darknet19.state_dict().items()):\n",
    "                if 'num_batches_tracked' in k:\n",
    "                    continue\n",
    "                param = torch.from_numpy(np.asarray(h5f[k]))\n",
    "                v.copy_(param)\n",
    "        self.darknet19.cuda()\n",
    "    def predict(self, fname, target=None):\n",
    "        self.darknet19.eval()\n",
    "        # 图片预处理\n",
    "        image = cv2.imread(fname)\n",
    "        image2 = cv2.resize(image, Config.inp_size) #(416, 416)\n",
    "        im = cv2.cvtColor(image2, cv2.COLOR_BGR2RGB) / 255.\n",
    "        im = np.expand_dims(im, 0)\n",
    "        im = Variable(torch.from_numpy(im).type(torch.FloatTensor), requires_grad=False).cuda().permute(0, 3, 1, 2)\n",
    "        print(\"Input size:\", im.shape)\n",
    "        # 预测\n",
    "        bbox_pred, iou_pred, prob_pred = self.darknet19(im)\n",
    "        bbox_pred = bbox_pred.data.cpu().numpy()\n",
    "        iou_pred = iou_pred.data.cpu().numpy()\n",
    "        prob_pred = prob_pred.data.cpu().numpy()\n",
    "        # 端处理\n",
    "        num_classes = voc.num_classes\n",
    "        anchors = voc.anchors\n",
    "        num_anchors = voc.num_anchors\n",
    "        W, H = Config.out_size\n",
    "        cell_size = 1. / W      #假设图片尺寸为1x1\n",
    "        assert bbox_pred.shape[0] == 1, 'postprocess only support one image per batch'\n",
    "        #将[x,y,w,h]变为[x1,y1,x2,y2]\n",
    "        bbox_out = np.zeros((1, H*W, voc.num_anchors, 4), dtype=np.float)\n",
    "        for row in range(H):\n",
    "            for col in range(W):\n",
    "                ind = row * W + col\n",
    "                for a in range(num_anchors):\n",
    "                    cx = (bbox_pred[0, ind, a, 0] + col) * cell_size\n",
    "                    cy = (bbox_pred[0, ind, a, 1] + row) * cell_size\n",
    "                    bw = bbox_pred[0, ind, a, 2] * anchors[a][0] * 0.5 * cell_size\n",
    "                    bh = bbox_pred[0, ind, a, 3] * anchors[a][1] * 0.5 * cell_size\n",
    "\n",
    "                    bbox_out[0, ind, a, 0] = cx - bw\n",
    "                    bbox_out[0, ind, a, 1] = cy - bh\n",
    "                    bbox_out[0, ind, a, 2] = cx + bw\n",
    "                    bbox_out[0, ind, a, 3] = cy + bh\n",
    "        \n",
    "        bbox_pred = np.reshape(bbox_out, [-1, 4])\n",
    "        bbox_pred[:, 0::2] *= float(Config.inp_size[1])\n",
    "        bbox_pred[:, 1::3] *= float(Config.inp_size[0])\n",
    "        # bbox_pred[:, 0::2] *= float(image.shape[1])\n",
    "        # bbox_pred[:, 1::3] *= float(image.shape[0])\n",
    "        bbox_pred = bbox_pred.astype(np.int)\n",
    "        iou_pred = np.reshape(iou_pred, [-1])\n",
    "        prob_pred = np.reshape(prob_pred, [-1, num_classes])\n",
    "        cls_inds = np.argmax(prob_pred, axis=1) #一维[]记录最大值的索引\n",
    "        prob_pred = prob_pred[(np.arange(prob_pred.shape[0]), cls_inds)]\n",
    "        scores = iou_pred * prob_pred\n",
    "        assert len(scores) == len(bbox_pred), '{}, {}'.format(scores.shape, bbox_pred.shape)\n",
    "        # 阈值\n",
    "        keep = np.where(scores >= Config.thresh)\n",
    "        bbox_pred = bbox_pred[keep]\n",
    "        scores = scores[keep]\n",
    "        cls_inds = cls_inds[keep]\n",
    "        #nms\n",
    "        keep = np.zeros(len(bbox_pred), dtype=np.int)\n",
    "        for i in range(voc.num_classes):\n",
    "            inds = np.where(cls_inds == i)[0]   #类别为i的bbox的索引\n",
    "            if len(inds) == 0:\n",
    "                continue\n",
    "            c_bboxes = bbox_pred[inds]\n",
    "            c_scores = scores[inds]\n",
    "            dets = np.hstack((c_bboxes, c_scores[:, np.newaxis])).astype(np.float32)#[x1,y1,x2,y2,scores]\n",
    "            c_keep = self.cls_nms(dets)\n",
    "            keep[inds[c_keep]] = 1\n",
    "        keep = np.where(keep > 0)[0]\n",
    "        bbox_pred = bbox_pred[keep]\n",
    "        scores = scores[keep]\n",
    "        cls_inds = cls_inds[keep]\n",
    "        # w和h经过换算之后可能使得(x1,y1)和(x2,y2)超出照片边界\n",
    "        # if bbox_pred.shape[0] != 0:\n",
    "        #     bbox_pred[:, 0] = np.maximum(np.minimum(bbox_pred[:, 0], Config.inp_size[1]-1), 0)\n",
    "        #     bbox_pred[:, 1] = np.maximum(np.minimum(bbox_pred[:, 1], Config.inp_size[0]-1), 0)\n",
    "        #     bbox_pred[:, 2] = np.maximum(np.minimum(bbox_pred[:, 2], Config.inp_size[1]-1), 0)\n",
    "        #     bbox_pred[:, 3] = np.maximum(np.minimum(bbox_pred[:, 3], Config.inp_size[0]-1), 0)\n",
    "        imgcv = np.copy(image2)\n",
    "        h, w, _ = imgcv.shape\n",
    "        bbox_pred = bbox_pred.astype(np.int)\n",
    "        print(bbox_pred)\n",
    "        for i, box in enumerate(bbox_pred):\n",
    "            cls_indx = cls_inds[i]\n",
    "            thick = int((h + w) / 300)\n",
    "            cv2.rectangle(imgcv,\n",
    "                          (box[0], box[1]), (box[2], box[3]),\n",
    "                          voc.colors[cls_indx], 2)\n",
    "            mess = '%s: %.3f' % (voc.label_names[cls_indx], scores[i])\n",
    "            cv2.putText(imgcv, mess, (box[0], box[1] - 12),\n",
    "                    0, 1e-3 * h, voc.colors[cls_indx], thick // 3)\n",
    "        cv2.imwrite('v2-result.jpg',imgcv)\n",
    "\n",
    "    def cls_nms(self, dets):\n",
    "        x1 = dets[:, 0]\n",
    "        y1 = dets[:, 1]\n",
    "        x2 = dets[:, 2]\n",
    "        y2 = dets[:, 3]\n",
    "        scores = dets[:, 4]\n",
    "        areas = (x2 - x1) * (y2 - y1)\n",
    "        order = scores.argsort()[::-1]#从大到小排序的索引\n",
    "        ndets = dets.shape[0]\n",
    "        suppressed = np.zeros((ndets), dtype=np.int)\n",
    "        keep = []\n",
    "        while order.size > 0:\n",
    "            i = order[0]\n",
    "            keep.append(i)\n",
    "            xx1 = np.maximum(x1[i], x1[order[1:]])\n",
    "            yy1 = np.maximum(y1[i], y1[order[1:]])\n",
    "            xx2 = np.minimum(x2[i], x2[order[1:]])\n",
    "            yy2 = np.minimum(y2[i], y2[order[1:]])\n",
    "\n",
    "            w = np.maximum(0.0, xx2 - xx1 + 1)\n",
    "            h = np.maximum(0.0, yy2 - yy1 + 1)\n",
    "            inter = w * h\n",
    "            ovr = inter / (areas[i] + areas[order[1:]] - inter)\n",
    "\n",
    "            inds = np.where(ovr <= Config.thresh)[0]\n",
    "            #ids = (ious <= Config.nms_thresh).nonzero()\n",
    "            if len(inds) == 0:\n",
    "                break\n",
    "            \n",
    "            order = order[inds+1]\n",
    "\n",
    "        return keep        "
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "source": [
    "i = cv2.imread('/home/yinzp/gitee/ObjectDetection-2D/image/person.jpg')\n",
    "print(i.shape)"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "(424, 640, 3)\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "source": [
    "yolov2 = YOLOV2(path_ckpt='/home/yinzp/gitee/ckpt/yolo-voc.weights.h5')\n",
    "yolov2.predict('/home/yinzp/gitee/ObjectDetection-2D/image/person.jpg')"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Use ckpt from: /home/yinzp/gitee/ckpt/yolo-voc.weights.h5\n",
      "Input size: torch.Size([1, 3, 416, 416])\n",
      "xx1 [257. 264.]\n",
      "ious [-0. -0.]\n",
      "order [2 1 0]\n",
      "ids [0 1]\n",
      "xx1 [264.]\n",
      "ious [-0.]\n",
      "order [1 0]\n",
      "ids [0]\n",
      "xx1 [119.]\n",
      "ious [-0.]\n",
      "order [1 0]\n",
      "ids [0]\n",
      "[[119 100 178   0]\n",
      " [264 115 393   0]\n",
      " [118  96 180   0]\n",
      " [257 135 377   0]\n",
      " [257 139 393   0]\n",
      " [ 51 257 124   0]]\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "/home/yinzp/.local/lib/python3.6/site-packages/torch/nn/functional.py:1805: UserWarning: nn.functional.sigmoid is deprecated. Use torch.sigmoid instead.\n",
      "  warnings.warn(\"nn.functional.sigmoid is deprecated. Use torch.sigmoid instead.\")\n",
      "/home/yinzp/.local/lib/python3.6/site-packages/ipykernel_launcher.py:70: UserWarning: Implicit dimension choice for softmax has been deprecated. Change the call to include dim=X as an argument.\n"
     ]
    }
   ],
   "metadata": {}
  }
 ],
 "metadata": {
  "orig_nbformat": 4,
  "language_info": {
   "name": "python",
   "version": "3.6.9",
   "mimetype": "text/x-python",
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "pygments_lexer": "ipython3",
   "nbconvert_exporter": "python",
   "file_extension": ".py"
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3.6.9 64-bit"
  },
  "interpreter": {
   "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}